#ifndef Pre_ParAssembly_h
#define Pre_ParAssembly_h

#include "Utility/Algebra/AlgebraModule.h"

#ifdef SG_USE_MPI

#include "Utility/Algebra/AlgebraicTypes.h"
#include "Utility/Algebra/Matrix.h"
#include "Utility/Algebra/Vector3D.h"

#include "DataStructure/Common/Id.h"
#include "DataStructure/Common/Status.h"
#include "DataStructure/Input/Load/LoadCombination.h"
#include "DataStructure/Input/Solver/FieldTypes.h"
#include "DataStructure/Input/Solver/Subcase.h"

#include "DBManager/DBServiceFactory.h"

#include "SGFem/Pre/Common/PreExport.h"
#include "SGFem/Pre/DOFHandler/DOFHandler.h"

SG_NAMESPACE_OPEN
namespace FEMSolver
{
    // alias to avoid the tiresome names.
    typedef SG::Algebra::Vec_t   Vec_t;
    typedef SG::Algebra::SpMat_t SpMat_t;

    typedef SG::Algebra::Complex_t ComplexVec_t;    

    // maitain the compatibilities with other modules
    using EigenSparse = Eigen::SparseMatrix<double>;
    using EigenVector = Eigen::VectorX<double>;

    /**
     * @brief Parallel Assembly to accumulate the contributons from all the entties within every partition
     *
     * In the conventional finite element method, the well-posed mathematical problems, which are
     * composed of the governing equations and inital/boundary conditions, firstly, will be discretized
     * over non-overlapping geometry elements. And then, the spatially discretized equations can be
     * formed by summing contributions from the entities (elements/nodes/conditions). These contributions
     * form the left-hand matrix term and/or the right-hand terms in the resultant equations (ODE/DAE).
     * At this stage, a semi-discretized equations can be built to approximately describe the original
     * physical model.
     *
     * Take the structural linear static analysis as an example, the semi-discretized equations can
     * be described as,
     * \f[
     *  \begin{pmatrix} \boldsymbol{K}_{ff} & \boldsymbol{K}_{fm}\\ \boldsymbol{K}_{mf} & \boldsymbol{K}_{mm}\end{pmatrix}
     *  \begin{pmatrix} \boldsymbol{u}_f\\ \boldsymbol{u}_m \end{pmatrix} =
     *  \begin{pmatrix} \boldsymbol{F}_{f}\\ \boldsymbol{F}_{m} \end{pmatrix}
     * \f]
     *
     * where, \f$ \boldsymbol{u}_f \f$ the active DOF (including the fixed DOF), \f$ \boldsymbol{u}_m \f$ the slave DOF.
     *
     * The additional MPC can be given as,
     * \f[
     *  \boldsymbol{R}_{n} \boldsymbol{u}_{f} = \boldsymbol{R}_{m} \boldsymbol{u}_{m}
     * \f]
     *
     * Let \f$ \boldsymbol{G}_{m} = \boldsymbol{R}_{m}^{-1} \boldsymbol{R}_{f}\f$, the MPC can be
     * re-written as,
     * \f[
     *  \boldsymbol{u}_{m} = \boldsymbol{G}_{m} \boldsymbol{u}_{f}
     * \f]
     *
     * Furtherly, large-scale numerical simulations on parallel computers require the distribution
     * of the ﬁnite element mesh to the processors. In this situations, the contributions from all
     * the entties within every partition should be summed up to form the global discretized equations.
     *
     * Thus, a new utility should be introduced to maitain the tasks mentioned above.
     *
     * @author nene
     * @date March, 2025.
     */
    class Pre_DECLSPEC ParAssembly
    {
      private:
        /**
         * @brief    Construct a new Assembly Eigen object
         *
         * m_MassType 赋值为 COUPMASS的默认值
         */
        ParAssembly () : m_isInitialized{ false }, m_needMass{ false }, m_MassType{ -1 }, m_DOFHandler ()
        {}
        ParAssembly (const ParAssembly&);
        ParAssembly& operator= (const ParAssembly&);
        ~ParAssembly () = default;

        /// @brief 生成矩阵方程的维度， 不包括MPC从自由度。
        std::size_t GetAcctiveDOFSize () const;

        /** @brief 计算mpc导致的约束方程， um = G *u_n */
        void ComputeMPCControlMatrix ();

        /** @brief 获取控制矩阵的组装结果
         *
         * 缓存这一计算结果，主要是在后续从自由位移恢复时使用
         */
        const SpMat_t& GetControlMatrix () const;

        /**  @brief eigen稀疏矩阵分块
         *
         * 按照 SplitIndex 将稀疏矩阵分割成四个小的子块
         */
        static void Matrixchunks (
            const SpMat_t& sparse, std::size_t splitIndex, SpMat_t& _OUT mat1, SpMat_t& _OUT mat2, SpMat_t& _OUT mat3, SpMat_t& _OUT mat4);

      public:
        /// singleton getter
        static ParAssembly& GetInstance ();

      public:
        /** @brief 初始化，完成节点和自由度排序。开辟整体矩阵及向量的存储空间
         *
         * 一组具有相同边界条件的多个工况，刚度矩阵的组装可以只进行一次，所有的右端项也可以同时求出。
         *
         * @param [in] dbServiceFactory   网格模型数据
         * @param [in] subcaseId0  工况id, 用于边界条件选择
         * @param [in] loadCases   除开Load（Case）之外的载荷集编号，用于开辟右端项储存空间.对模态分析不需要右端项使用空默认值
         */
        void Initialize (const SG::DBManager::DBServiceFactorySP        dbServiceFactory,
                         const SG::DataStructure::FEM::Subcase&         subcase,
                         const std::set<SG::DataStructure::Common::Id>& loadCases = {});

        /** 初始化是否执行
         *
         * 避免重复调用初始化函数 Initialize。
         * */
        bool IsInitialized () const;

        /** @brief 清空已经累加的元素，如果更新计算模型或者计算不同的边界条件，需要在调用Initialize之前，调用本函数。
         */
        void Clear ();

        /** @brief 查询节点对应的索引*/
        SG::DataStructure::Common::Status FindNodeIndex (SG::DataStructure::Common::Id nodeId, _OUT std::size_t& nodeIndex) const;

        /** @brief 查询节点数据在数据库存储的索引位置
         * @param [in] nodeIds      节点编号集合， 可传入std::vector或者数组 的首地址
         * @param [out] nodeIndexes 节点索引集合
         * @attention  确保 nodeIndexes和nodeIds具有相同的大小（size）
         */
        void FindNodeIndex (const SG::DataStructure::Common::Id* nodeIds, _OUT std::vector<std::size_t>& nodeIndexes) const;

        /** @brief 获得指定索引节点对应自由度的全局编号
         *
         * @param [in] nodeIndexes  节点集索引
         * @param [in] DOFs         指定的自由度集
         * @param [out] dofGlobalId 返回自由度全局编号
         */
        void GetGlobalDOFId (const std::vector<std::size_t>&       nodeIndexes,
                             const SG::DataStructure::FEM::DofSet& DOFs,
                             _OUT std::vector<int>& dofGlobalIds) const;

        /** 查询节点某个方向上的自由度
         *
         * @param nodeIndex 节点对应的索引
         */
        void GetGlobalDOFId (const std::size_t nodeIndex, const SG::DataStructure::FEM::DOFDirections& dof, _OUT int& dofGlobalId) const;

        /**  @brief 查询节点指定自由度的全局编号
         *
         * @param [in] nodeId 节点编号
         * @param [in] DOFs   指定的自由度集
         * @param [out] dofGlobalId 返回自由度全局编号
         */
        void ComputeGlobalDOFId (const SG::DataStructure::Common::Id   nodeId,
                                 const SG::DataStructure::FEM::DofSet& DOFs,
                                 _OUT std::vector<int>& dofGlobalIds);

        /** @name 叠加单元级别贡献
         */
        /// @{
        /** @brief 单元刚度矩阵添加
         * @param globalIndex 单元矩阵对应的全局自由度
         * @param stiffness   单元刚度矩阵（位移输出坐标系下）
         */
        void AccumulateStiffness (const std::vector<int>& globalIndex, const SG::Algebra::Matrix& stiffness);

        /** @brief 单元刚度矩阵添加
         * @param globalIndex 单元矩阵对应的全局自由度
         * @param mass        单元协调质量矩阵（位移输出坐标系下）
         */
        void AccumulateCoupledMass (const std::vector<int>& globalIndex, const SG::Algebra::Matrix& mass);

        /** @brief 累加应力刚度的贡献 */
        void AccumulateStressStiffness (const std::vector<int>& globalIndex, const SG::Algebra::Matrix& kSigma);

        /** @brief 组装对载荷右端项贡献
         *
         * 计入集中载荷对于右端项的贡献
         */
        void AccumulateLoad (SG::DataStructure::Common::Id sid, const std::vector<int>& globalDof, const Algebra::Vector3D& load);

        /** @brief 组装对载荷右端项贡献
         *
         * 计入集中载荷对于右端项的贡献
         */
        void AccumulateLoad (SG::DataStructure::Common::Id sid, const std::vector<int>& globalDof, const SG::Algebra::Matrix& load);

        /** @brief 累加控制矩阵
         *
         * 累计R单元和MPC对控制矩阵的贡献
         */
        void AccumulationControlMatrix (const std::vector<int>& gdofn, const std::vector<int>& gdofm, const SG::Algebra::Matrix& RnRm);
        /// @}

        /** @name 边界条件产生的特殊载荷 （针对热分析添加）
         */
        /// @{
        /** @ 初始化时调用此接口，
         * 调用此接口， 相应边界条件记录编号会被清空
         */
        void ClearBCLoad ();

        /** @brief 组装特别边界条件对载荷右端项贡献
         *
         * 此时 sid 的数据会被记录，
         * 之后返回右端项时，会叠加此sid 对应载荷的贡献
         */
        void AccumulateBCLoad (SG::DataStructure::Common::Id sid, const std::vector<int>& globalDof, const SG::Algebra::Matrix& load);
        /// @}

        /** @brief 是否需要计算控制矩阵
         *
         */
        bool NeedComputeControlMatrix () const;

        /** @brief 约束处理， 消去mpc的非独立自由度， 处理spc约束，添加autospc计算功能
         */
        void Assembly (const SG::DBManager::DBServiceFactorySP dbServiceFactory, const SG::DataStructure::Common::Id subcaseId);

        /// @brief 计算非Load组合对应的载荷右端项
        void ComputeBVector (SG::DataStructure::Common::Id loadCase, _OUT Vec_t& b) const;

        /// @brief 计算载荷组合(Load)的情况下载荷右端向量
        void ComputeBVector (const SG::DataStructure::FEM::LoadCombination& load, _OUT Vec_t& b) const;

        /**
         * @brief   : 针对动力学有关的分析， 需要激活对质量矩阵的约束处理
         * @author  : wanghuan
         * @date    : 2024.11.13
         */
        void ActiveMassMatrix (int massType = -1);

        /**
         * @brief    质量矩阵类型(`集中` & `耦合`)
         *
         * @return
         * @return
         * @author   wanghuan
         * @date     2024.11.26
         */
        bool NeedCoupleMass () const;

        /** @brief 获取完成约束处理的刚度矩阵的刚度矩阵
         *
         */
        const SpMat_t& GetStiffnessMatrix () const;

        /** @brief 获取结构一致质量矩阵
         *
         */
        const SpMat_t& GetCoupledMassMatrix () const;

        /** @brief 对应力刚度进行约束处理
         *
         */
        const SpMat_t ComputeStressStiffnessMatrix ();

        /** @brief 位移恢复
         * @param[in] displacement 矩阵计算得到的位移向量, 变量会被修改
         * @param[out] extendDisplacement 位移恢复延展后的位移向量
         */
        template <typename DVector>
        Pre_DECLSPEC void RecoveryDisplacement (DVector& displacement, _OUT DVector& extendDisplacemen) const;

        /// @brief 包含所有约束的总自由度数量
        std::size_t GetTotalDOFSize () const;

        /**
         * @brief 获取处理完约束的自由度数量
         *
         * 用于动力学初始化 速度， 位移及 加速度矢量
         *
         * @return
         * @author   wanghuan
         * @date     2024.12.25
         */
        std::size_t GetFreeDOFSize () const;

        /** 获取所有节点id */
        const std::vector<SG::DataStructure::Common::Id>& GetNodeIds () const;

      public:
        bool                                m_isInitialized;           ///< 类是否被初始化
        SG::DataStructure::FEM::Field_Types m_fieldType;
        bool                                m_needMass;
        int                                 m_MassType;                ///< 当 其数值小于0，生成集中质量矩阵
        SG::FEMSolver::DOFHandler           m_DOFHandler;              ///< 自由度处理类， 用于提供节点自由度的全局索引

        std::map<SG::DataStructure::Common::Id, std::size_t> m_loadmap;
        std::vector<Vec_t>                                   m_FnSet;  ///< 自由自由度载荷向量
        std::vector<Vec_t>                                   m_FmSet;  ///< MPC约束自由度载荷向量
        /** nSet 分割成 spc对应的自由度集合sSet 和 独立的自由度集合fSet */
        std::vector<Vec_t> m_FfSet;  ///< 独立自由度载荷向量

        /** @name 边界条件产生的特殊载荷 （针对热分析添加）
         */
        /// @{
        /**
         * @brief    边界条件产生的载荷
         *
         */
        std::map<SG::DataStructure::Common::Id, Vec_t> m_FnBCSet;      ///< 特殊边界条件产生的自由自由度载荷向量
        std::map<SG::DataStructure::Common::Id, Vec_t> m_FmBCSet;      ///< 特殊边界条件产生的MPC约束自由度载荷向量
        std::map<SG::DataStructure::Common::Id, Vec_t> m_FfBCSet;      ///< 特殊边界条件产生的独立自由度载荷向量
        std::set<SG::DataStructure::Common::Id>        m_BCLoadIdSet;  ///< 特别边界条件载荷sid，
        /// @}

        SpMat_t m_TriKff;         ///< 自由的自由度对应的刚度矩阵
        SpMat_t m_TriKfm;         ///< 自由自由度和多点约束控制自由度交叉部分
        SpMat_t m_TriKmm;         ///< 多点约束控制自由度对应刚度矩阵

        SpMat_t m_TriCoupledMff;  ///< 自由的自由度对应的       质量矩阵
        SpMat_t m_TriCoupledMmm;  ///< 多点约束控制自由度对应的 质量矩阵
        SpMat_t m_TriCoupledMfm;  ///< 交叉部分对应的           质量矩阵

        SpMat_t m_TriK4ff;        ///< 自由的自由度对应的       应力刚度矩阵
        SpMat_t m_TriK4mm;        ///< 多点约束控制自由度对应的 应力刚度矩阵
        SpMat_t m_TriK4fm;        ///< 交叉部分对应的           应力刚度矩阵

        SpMat_t m_RnTri;          ///< 用于控制矩阵计算
        SpMat_t m_RmTri;

        SpMat_t m_Gm;             ///< 通过主自由度获得从自由度的矩阵
        SpMat_t m_K;              ///< 刚度矩阵
        SpMat_t m_CoupledM;       ///< 单元协调质量矩阵矩阵

        SpMat_t m_kfs;            ///< 分块矩阵, 用于后续修正载荷右端项
        Vec_t m_Ys;             ///< SPC对应的位移值。用于后续位移恢复填充
    };

    inline bool ParAssembly::IsInitialized () const
    {
        return m_isInitialized;
    }

    /**清空上一次的刚度矩阵，载荷向量等数据
     * 不应该修改和网格有关的数据，如坐标系
     */
    inline void ParAssembly::Clear ()
    {
        m_isInitialized = false;
        m_needMass      = false;
        m_MassType      = -1;
        m_TriKff.clear ();
        m_TriKfm.clear ();
        m_TriKmm.clear ();

        m_RnTri.clear ();
        m_RmTri.clear ();

        m_FnSet.clear ();
        m_FmSet.clear ();

        m_TriCoupledMff.clear ();
        m_TriCoupledMmm.clear ();
        m_TriCoupledMfm.clear ();

        m_TriK4ff.clear ();
        m_TriK4mm.clear ();
        m_TriK4fm.clear ();
        // m_Constrain.clear();

        ClearBCLoad ();
    }


    /** @brief 是否需要计算控制矩阵
     *
     * 当网格中具有从节点就需要计算控制矩阵
     */
    inline bool ParAssembly::NeedComputeControlMatrix () const
    {
        return (m_DOFHandler.GetSlaveDOFSize () > 0);
    }

    inline const SpMat_t& ParAssembly::GetControlMatrix () const
    {
        return m_Gm;
    }

    inline void ParAssembly::ActiveMassMatrix (int massType)
    {
        m_needMass = true;
        m_MassType = massType;
    }

    inline bool ParAssembly::NeedCoupleMass () const
    {
        return m_MassType > 0;
    }

    inline const SpMat_t& ParAssembly::GetStiffnessMatrix () const
    {
        return m_K;
    }

    inline const SpMat_t& ParAssembly::GetCoupledMassMatrix () const
    {
        return m_CoupledM;
    }

    inline std::size_t ParAssembly::GetFreeDOFSize () const
    {
        return m_DOFHandler.ComputeSparseMatrixSize ();
    }

    inline std::size_t ParAssembly::GetAcctiveDOFSize () const
    {
        return m_DOFHandler.ComputeSparseMatrixSize () + m_DOFHandler.GetFixedDOFSize ();
    }

    inline std::size_t ParAssembly::GetTotalDOFSize () const
    {
        return GetAcctiveDOFSize () + m_DOFHandler.GetSlaveDOFSize ();
    }

    inline const std::vector<SG::DataStructure::Common::Id>& ParAssembly::GetNodeIds () const
    {
        return m_DOFHandler.GetNodeIds ();
    }    
}  // namespace FEMSolver
SG_NAMESPACE_CLOSE

#endif  // SG_USE_MPI

#endif  // Pre_ParAssembly_h